package com.agentframework.core.agent;

import com.agentframework.core.action.AbstractAction;
import com.agentframework.core.action.Action;
import com.agentframework.core.environment.AbstractEnvironment;
import com.agentframework.core.exception.AgentException;
import com.agentframework.core.provider.HumanLlm;
import com.agentframework.core.provider.Llm;
import com.agentframework.core.provider.LlmType;
import jakarta.annotation.PreDestroy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.lang.NonNull;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 抽象智能体基类
 * @author wangjun
 **/
public abstract class AbstractAgent implements Agent, ApplicationContextAware, InitializingBean {

    protected static final Logger logger = LoggerFactory.getLogger(AbstractAgent.class);

    public final String PREFIX_TEMPLATE = """
            名字 {name}, {profile}.
            """;

    protected final String agentId;

    protected String agentName;

    protected String agentProfile;

    protected String goal = "";

    protected List<Action> actions = new ArrayList<>();

    protected final AgentContext context;

    protected Llm llm;

    protected LlmType llmType;

    protected boolean isHuman;

    protected final List<String> states = new ArrayList<>();

    protected ReactStrategy reactStrategy;

    protected ApplicationContext applicationContext;

    protected List<ToolCallback> toolCallbacks = List.of();

    protected Map<String, Object> toolContext = Map.of();

    private ExecutorService messageProcessorExecutorService;

    protected AbstractAgent(String agentName, LlmType llmType) {
        this.agentName = agentName;
        this.llmType = llmType;
        this.agentId = UUID.randomUUID().toString().substring(0, 8);
        this.context = new AgentContext();
    }

    protected AbstractAgent(String agentName) {
        this(agentName, LlmType.DEEPSEEK);
    }

    @Override
    public void run(String message) {

        messageProcessorExecutorService.execute(() -> {
            while (true) {
                try {
                    AgentHolder.setAgent(this);
                    if (StringUtils.hasText(message)) {
                        logger.debug("agent {} 处理消息 : {}", agentName, message);
                        context.putMessage(UserMessage.builder().text(message).build());
                    }

                    if (observe() == 0) {
                        logger.debug("{}, 没有有效的消息，继续下一次监听", this.agentName);
                        continue;
                    }

                    AssistantMessage response = react();
                    if (response != null && StringUtils.hasText(response.getText())) {
                        var sendTo = response.getMetadata().get(SEND_TO);
                        if (sendTo == null || ((Set<?>) sendTo).isEmpty()) {
                            logger.debug("agent {} response message not set sendTo {}", this, response);
                            continue;
                        }
                        response.getMetadata().put(MSG_FROM, this.getAgentId());
                        publishMessage(response);
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                } catch (AgentException e) {
                    logger.error("agent {} - {} error {}", getAgentName(), getAgentId(), e.getMessage());
                    processActException();
                } finally {
                    AgentHolder.clear();
                }
            }
        });

    }

    @Override
    public void run() {
        run(null);
    }

    protected int observe() throws InterruptedException {
        // 清空上一次处理的消息
        context.getNews().clear();

        // 从队列中获取全部消息
        Message message = context.takeMessage();
        List<Message> newsDeque = context.popAllMessage();
        if (message != null) {
            newsDeque.add(message);
        }

        logger.debug("agent {} deque messages {}", this, newsDeque);

        // 加载需要处理的历史消息
        Set<Message> history = getHistoryMessages().stream().filter(
                m -> {
                    Object msgFrom = m.getMetadata().get(MSG_FROM);
                    Object sendTo = m.getMetadata().get(SEND_TO);
                    return (msgFrom != null && context.getWatch().contains(msgFrom.toString()))
                            || (sendTo != null && ((Set<?>) sendTo).contains(this.agentId));
                })
                .collect(Collectors.toSet());

        List<Message> news = newsDeque.stream().filter(n -> !history.contains(n)).toList();

        // 找出指定智能体只读消息
        List<Message> viewNews = news.stream().filter(n -> {
            Object viewTo = n.getMetadata().get(VIEW_TO);
            return viewTo != null && (((Set<?>) viewTo).contains(TO_ALL) || ((Set<?>) viewTo).contains(agentId));
        }).toList();

        // 只读消息直接加到记忆即可
        addToHistoryMessages(viewNews);

        // 新的消息需要排除掉只读的消息，因为不需要丢到智能体队列进行消息处理
        HashSet<Message> viewNewsSet = new HashSet<>(viewNews);
        news = news.stream().filter(n -> !viewNewsSet.contains(n)).toList();
        context.getNews().addAll(news);

        logger.debug("{} news {}", this, context.getNews());

        return context.getNews().size();
    }

    public void clear(){
        llm.getChatMemory().clear(getEnvId());
    }

    protected AssistantMessage react() {

        try {
            if (actions.isEmpty()) {
                return act();
            }

            AssistantMessage response = null;
            int actionTaken = 0;
            while (actionTaken < context.getMaxReactLoop()) {
                logger.debug("actionToken {}", actionTaken);
                boolean todo = think();
                if (!todo) {
                    break;
                }

                response = act();

                actionTaken ++;
            }

            setActionIndex(-1);

            return response;
        } catch (Exception e) {
            throw new AgentException(e);
        }
    }

    protected boolean think() {
        return reactStrategy.react(this);
    }

    protected AssistantMessage act() {
        if (context.getTodo() != null) {
            return context.getTodo().execute(context.getNews());
        } else {
            return (AssistantMessage) llm.call(getLatestMessage().getText());
        }
    }

    protected void processActException() {
        logger.error("{} error process message {}", this, context.getNews().get(0));
        Message errorMessage = context.getErrorMessage();
        if (errorMessage != null) {
            logger.error("{} error response message {}", this, errorMessage);
        }
    }

    public void addAction(Action action) {
        addActions(List.of(action));
    }

    public void addActions(List<Action> actions) {
        this.actions.addAll(actions);
        for(int i = 0; i< this.actions.size(); i++) {
            this.states.add(i + "." + ((AbstractAction) this.actions.get(i)).getActionName());
        }
    }

    public void setActions(List<Action> agentActions) {
        addActions(agentActions);
    }

    protected void setActionIndex(int index) {
        context.setActionIndex(actions.isEmpty() ? -1 : index);
        context.setTodo(!actions.isEmpty() && index >= 0 && index < actions.size() ? actions.get(index) : null);
    }

    @Override
    public AgentContext getContext() {
        return this.context;
    }

    @Override
    public String getAgentId() {
        return this.agentId;
    }

    @Override
    public List<Action> getActions() {
        return actions;
    }

    public LlmType getLlmType() {
        return this.llmType;
    }

    public void setLlmType(LlmType llmType) {
        this.llmType = llmType;
    }

    public List<Message> getHistoryMessages() {
        return llm.getChatMemory().get(getEnvId());
    }

    public void addToHistoryMessages(List<Message> messages) {
        if (!messages.isEmpty()) {
            llm.getChatMemory().add(getEnvId(), messages);
        }
    }

    public String getEnvId() {
        return getContext().getEnvironment() == null ? getAgentId() : getContext().getEnvironment().getEnvId();
    }

    /**
     * 获取队列中的最新一条消息
     * @return message
     */
    protected Message getLatestMessage() {
        List<Message> messages = context.getNews();
        return messages.isEmpty() ? null : messages.get(0);
    }

    protected void setReactModel(Class<? extends ReactStrategy> agentReactMode) {
        this.context.setReactMode(agentReactMode);
    }

    public List<ToolCallback> getToolCallbacks() {
        return toolCallbacks;
    }

    public void setToolCallbacks(Object ... toolCallback) {
        this.toolCallbacks = Arrays.stream(ToolCallbacks.from(toolCallback)).toList();
    }

    public void setToolContext(Map<String, Object> toolContext) {
        this.toolContext = toolContext;
    }

    public Map<String, Object> getToolContext() {
        return this.toolContext;
    }

    public String getAgentName() {
        return agentName;
    }

    public void setAgentName(String agentName) {
        this.agentName = agentName;
    }

    public String getAgentProfile() {
        return agentProfile;
    }

    public String getGoal() {
        return goal;
    }

    public void setGoal(String goal) {
        this.goal = goal;
    }

    @Override
    public AbstractEnvironment getEnvironment() {
        return getContext().getEnvironment();
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        this.llm = applicationContext.getBean(Llm.class);
        this.reactStrategy = applicationContext.getBean(ReactStrategy.class);
        if (isHuman) {
            this.llm = applicationContext.getBean(HumanLlm.class);
        }
        this.messageProcessorExecutorService = Executors.newSingleThreadExecutor(r -> {
            Thread t = new Thread(r);
            t.setName("AgentProcessor-" + this.agentId);
            return t;
        });
    }

    @PreDestroy
    public void shutdown() {
        if (messageProcessorExecutorService != null && !messageProcessorExecutorService.isShutdown()) {
            messageProcessorExecutorService.shutdownNow();
            try {
                if (!messageProcessorExecutorService.awaitTermination(1, TimeUnit.SECONDS)) {
                    logger.warn("Agent {} 消息处理器关闭超时", agentName);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }

    @Override
    public void publishMessage(Message message) {
        context.getEnvironment().publishMessage(message);
    }

    @Override
    public String toString() {
        return "Agent{" +
                "agentName='" + agentName + '\'' +
                ", agentId='" + agentId + '\'' +
                '}';
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        AbstractAgent that = (AbstractAgent) o;
        return Objects.equals(agentId, that.agentId);
    }

    @Override
    public int hashCode() {
        return Objects.hash(agentId);
    }

}
